{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import json\n",
    "import re\n",
    "import random\n",
    "import numpy as np\n",
    "from collections import defaultdict\n",
    "\n",
    "import cn2an\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from nl2sql.utils import read_data, read_tables, SQL, Query, Question, Table\n",
    "from keras_bert import get_checkpoint_paths, load_vocabulary, Tokenizer, load_trained_model_from_checkpoint\n",
    "from keras.utils.data_utils import Sequence\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.layers import Input, Lambda, Dense\n",
    "from keras.models import Model\n",
    "from keras.optimizers import Adam\n",
    "from keras.utils import multi_gpu_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_table_file = '../data/train.tables.json'\n",
    "train_data_file = '../data/train.json'\n",
    "\n",
    "val_table_file = '../data/val.tables.json'\n",
    "val_data_file = '../data/val.json'\n",
    "\n",
    "test_table_file = '../data/test.tables.json'\n",
    "test_data_file = '../data/test.json'\n",
    "\n",
    "# Download pretrained BERT model from https://github.com/ymcui/Chinese-BERT-wwm\n",
    "bert_model_path = '../model/chinese_wwm_L-12_H-768_A-12'\n",
    "\n",
    "paths = get_checkpoint_paths(bert_model_path)\n",
    "\n",
    "task1_file = '../submit/task1_output.json'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tables = read_tables(train_table_file)\n",
    "train_data = read_data(train_data_file, train_tables)\n",
    "\n",
    "val_tables = read_tables(val_table_file)\n",
    "val_data = read_data(val_data_file, val_tables)\n",
    "\n",
    "test_tables = read_tables(test_table_file)\n",
    "test_data = read_data(test_data_file, test_tables)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_float(value):\n",
    "    try:\n",
    "        float(value)\n",
    "        return True\n",
    "    except ValueError:\n",
    "        return False\n",
    "\n",
    "def cn_to_an(string):\n",
    "    try:\n",
    "        return str(cn2an.cn2an(string, 'normal'))\n",
    "    except ValueError:\n",
    "        return string\n",
    "\n",
    "def an_to_cn(string):\n",
    "    try:\n",
    "        return str(cn2an.an2cn(string))\n",
    "    except ValueError:\n",
    "        return string\n",
    "\n",
    "def str_to_num(string):\n",
    "    try:\n",
    "        float_val = float(cn_to_an(string))\n",
    "        if int(float_val) == float_val:   \n",
    "            return str(int(float_val))\n",
    "        else:\n",
    "            return str(float_val)\n",
    "    except ValueError:\n",
    "        return None\n",
    "\n",
    "def str_to_year(string):\n",
    "    year = string.replace('年', '')\n",
    "    year = cn_to_an(year)\n",
    "    if is_float(year) and float(year) < 1900:\n",
    "        year = int(year) + 2000\n",
    "        return str(year)\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "def load_json(json_file):\n",
    "    result = []\n",
    "    if json_file:\n",
    "        with open(json_file) as file:\n",
    "            for line in file:\n",
    "                result.append(json.loads(line))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "class QuestionCondPair:\n",
    "    def __init__(self, query_id, question, cond_text, cond_sql, label):\n",
    "        self.query_id = query_id\n",
    "        self.question = question\n",
    "        self.cond_text = cond_text\n",
    "        self.cond_sql = cond_sql\n",
    "        self.label = label\n",
    "\n",
    "    def __repr__(self):\n",
    "        repr_str = ''\n",
    "        repr_str += 'query_id: {}\\n'.format(self.query_id)\n",
    "        repr_str += 'question: {}\\n'.format(self.question)\n",
    "        repr_str += 'cond_text: {}\\n'.format(self.cond_text)\n",
    "        repr_str += 'cond_sql: {}\\n'.format(self.cond_sql)\n",
    "        repr_str += 'label: {}\\n'.format(self.label)\n",
    "        return repr_str\n",
    "\n",
    "    \n",
    "class NegativeSampler:\n",
    "    \"\"\"\n",
    "    从 question - cond pairs 中采样\n",
    "    \"\"\"\n",
    "    def __init__(self, neg_sample_ratio=10):\n",
    "        self.neg_sample_ratio = neg_sample_ratio\n",
    "    \n",
    "    def sample(self, data):\n",
    "        positive_data = [d for d in data if d.label == 1]\n",
    "        negative_data = [d for d in data if d.label == 0]\n",
    "        negative_sample = random.sample(negative_data, \n",
    "                                        len(positive_data) * self.neg_sample_ratio)\n",
    "        return positive_data + negative_sample\n",
    "\n",
    "    \n",
    "class FullSampler:\n",
    "    \"\"\"\n",
    "    不抽样，返回所有的 pairs\n",
    "    \n",
    "    \"\"\"\n",
    "    def sample(self, data):\n",
    "        return data\n",
    "\n",
    "class CandidateCondsExtractor:\n",
    "    \"\"\"\n",
    "    params:\n",
    "        - share_candidates: 在同 table 同 column 中共享 real 型 candidates\n",
    "    \"\"\"\n",
    "    CN_NUM = '〇一二三四五六七八九零壹贰叁肆伍陆柒捌玖貮两'\n",
    "    CN_UNIT = '十拾百佰千仟万萬亿億兆点'\n",
    "    \n",
    "    def __init__(self, share_candidates=True):\n",
    "        self.share_candidates = share_candidates\n",
    "        self._cached = False\n",
    "    \n",
    "    def build_candidate_cache(self, queries):\n",
    "        self.cache = defaultdict(set)\n",
    "        print('building candidate cache')\n",
    "        for query_id, query in tqdm(enumerate(queries), total=len(queries)):\n",
    "            value_in_question = self.extract_values_from_text(query.question.text)\n",
    "            \n",
    "            for col_id, (col_name, col_type) in enumerate(query.table.header):\n",
    "                value_in_column = self.extract_values_from_column(query, col_id)\n",
    "                if col_type == 'text':\n",
    "                    cond_values = value_in_column\n",
    "                elif col_type == 'real':\n",
    "                    if len(value_in_column) == 1: \n",
    "                        cond_values = value_in_column + value_in_question\n",
    "                    else:\n",
    "                        cond_values = value_in_question\n",
    "                cache_key = self.get_cache_key(query_id, query, col_id)\n",
    "                self.cache[cache_key].update(cond_values)\n",
    "        self._cached = True\n",
    "    \n",
    "    def get_cache_key(self, query_id, query, col_id):\n",
    "        if self.share_candidates:\n",
    "            return (query.table.id, col_id)\n",
    "        else:\n",
    "            return (query_id, query.table.id, col_id)\n",
    "        \n",
    "    def extract_year_from_text(self, text):\n",
    "        values = []\n",
    "        num_year_texts = re.findall(r'[0-9][0-9]年', text)\n",
    "        values += ['20{}'.format(text[:-1]) for text in num_year_texts]\n",
    "        cn_year_texts = re.findall(r'[{}][{}]年'.format(self.CN_NUM, self.CN_NUM), text)\n",
    "        cn_year_values = [str_to_year(text) for text in cn_year_texts]\n",
    "        values += [value for value in cn_year_values if value is not None]\n",
    "        return values\n",
    "    \n",
    "    def extract_num_from_text(self, text):\n",
    "        values = []\n",
    "        num_values = re.findall(r'[-+]?[0-9]*\\.?[0-9]+', text)\n",
    "        values += num_values\n",
    "        \n",
    "        cn_num_unit = self.CN_NUM + self.CN_UNIT\n",
    "        cn_num_texts = re.findall(r'[{}]*\\.?[{}]+'.format(cn_num_unit, cn_num_unit), text)\n",
    "        cn_num_values = [str_to_num(text) for text in cn_num_texts]\n",
    "        values += [value for value in cn_num_values if value is not None]\n",
    "    \n",
    "        cn_num_mix = re.findall(r'[0-9]*\\.?[{}]+'.format(self.CN_UNIT), text)\n",
    "        for word in cn_num_mix:\n",
    "            num = re.findall(r'[-+]?[0-9]*\\.?[0-9]+', word)\n",
    "            for n in num:\n",
    "                word = word.replace(n, an_to_cn(n))\n",
    "            str_num = str_to_num(word)\n",
    "            if str_num is not None:\n",
    "                values.append(str_num)\n",
    "        return values\n",
    "    \n",
    "    def extract_values_from_text(self, text):\n",
    "        values = []\n",
    "        values += self.extract_year_from_text(text)\n",
    "        values += self.extract_num_from_text(text)\n",
    "        return list(set(values))\n",
    "   \n",
    "    def extract_values_from_column(self, query, col_ids):\n",
    "        question = query.question.text\n",
    "        question_chars = set(query.question.text)\n",
    "        unique_col_values = set(query.table.df.iloc[:, col_ids].astype(str))\n",
    "        select_col_values = [v for v in unique_col_values \n",
    "                             if (question_chars & set(v))]\n",
    "        return select_col_values\n",
    "    \n",
    "    \n",
    "class QuestionCondPairsDataset:\n",
    "    \"\"\"\n",
    "    question - cond pairs 数据集\n",
    "    \"\"\"\n",
    "    OP_PATTERN = {\n",
    "        'real':\n",
    "        [\n",
    "            {'cond_op_idx': 0, 'pattern': '{col_name}大于{value}'},\n",
    "            {'cond_op_idx': 1, 'pattern': '{col_name}小于{value}'},\n",
    "            {'cond_op_idx': 2, 'pattern': '{col_name}是{value}'}\n",
    "        ],\n",
    "        'text':\n",
    "        [\n",
    "            {'cond_op_idx': 2, 'pattern': '{col_name}是{value}'}\n",
    "        ]\n",
    "    }    \n",
    "    \n",
    "    def __init__(self, queries, candidate_extractor, has_label=True, model_1_outputs=None):\n",
    "        self.candidate_extractor = candidate_extractor\n",
    "        self.has_label = has_label\n",
    "        self.model_1_outputs = model_1_outputs\n",
    "        self.data = self.build_dataset(queries)\n",
    "        \n",
    "    def build_dataset(self, queries):\n",
    "        if not self.candidate_extractor._cached:\n",
    "            self.candidate_extractor.build_candidate_cache(queries)\n",
    "            \n",
    "        pair_data = []\n",
    "        for query_id, query in enumerate(queries):\n",
    "            select_col_id = self.get_select_col_id(query_id, query)\n",
    "            for col_id, (col_name, col_type) in enumerate(query.table.header):\n",
    "                if col_id not in select_col_id:\n",
    "                    continue\n",
    "                    \n",
    "                cache_key = self.candidate_extractor.get_cache_key(query_id, query, col_id)\n",
    "                values = self.candidate_extractor.cache.get(cache_key, [])\n",
    "                pattern = self.OP_PATTERN.get(col_type, [])\n",
    "                pairs = self.generate_pairs(query_id, query, col_id, col_name, \n",
    "                                               values, pattern)\n",
    "                pair_data += pairs\n",
    "        return pair_data\n",
    "    \n",
    "    def get_select_col_id(self, query_id, query):\n",
    "        if self.model_1_outputs:\n",
    "            select_col_id = [cond_col for cond_col, *_ in self.model_1_outputs[query_id]['conds']]\n",
    "        elif self.has_label:\n",
    "            select_col_id = [cond_col for cond_col, *_ in query.sql.conds]\n",
    "        else:\n",
    "            select_col_id = list(range(len(query.table.header)))\n",
    "        return select_col_id\n",
    "            \n",
    "    def generate_pairs(self, query_id, query, col_id, col_name, values, op_patterns):\n",
    "        pairs = []\n",
    "        for value in values:\n",
    "            for op_pattern in op_patterns:\n",
    "                cond = op_pattern['pattern'].format(col_name=col_name, value=value)\n",
    "                cond_sql = (col_id, op_pattern['cond_op_idx'], value)\n",
    "                real_sql = {}\n",
    "                if self.has_label:\n",
    "                    real_sql = {tuple(c) for c in query.sql.conds}\n",
    "                label = 1 if cond_sql in real_sql else 0\n",
    "                pair = QuestionCondPair(query_id, query.question.text,\n",
    "                                        cond, cond_sql, label)\n",
    "                pairs.append(pair)\n",
    "        return pairs\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_result = load_json(task1_file)\n",
    "\n",
    "tr_qc_pairs = QuestionCondPairsDataset(train_data, \n",
    "                                       candidate_extractor=CandidateCondsExtractor(share_candidates=False))\n",
    "\n",
    "te_qc_pairs = QuestionCondPairsDataset(test_data, \n",
    "                                       candidate_extractor=CandidateCondsExtractor(share_candidates=True),\n",
    "                                       has_label=False,\n",
    "                                       model_1_outputs=task1_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "class SimpleTokenizer(Tokenizer):\n",
    "    def _tokenize(self, text):\n",
    "        R = []\n",
    "        for c in text:\n",
    "            if c in self._token_dict:\n",
    "                R.append(c)\n",
    "            elif self._is_space(c):\n",
    "                R.append('[unused1]')\n",
    "            else:\n",
    "                R.append('[UNK]')\n",
    "        return R\n",
    "\n",
    "            \n",
    "def construct_model(paths, use_multi_gpus=False):\n",
    "    token_dict = load_vocabulary(paths.vocab)\n",
    "    tokenizer = SimpleTokenizer(token_dict)\n",
    "\n",
    "    bert_model = load_trained_model_from_checkpoint(\n",
    "        paths.config, paths.checkpoint, seq_len=None)\n",
    "    for l in bert_model.layers:\n",
    "        l.trainable = True\n",
    "\n",
    "    x1_in = Input(shape=(None,), name='input_x1', dtype='int32')\n",
    "    x2_in = Input(shape=(None,), name='input_x2')\n",
    "    x = bert_model([x1_in, x2_in])\n",
    "    x_cls = Lambda(lambda x: x[:, 0])(x)\n",
    "    y_pred = Dense(1, activation='sigmoid', name='output_similarity')(x_cls)\n",
    "\n",
    "    model = Model([x1_in, x2_in], y_pred)\n",
    "    if use_multi_gpus:\n",
    "        print('using multi-gpus')\n",
    "        model = multi_gpu_model(model, gpus=2)\n",
    "\n",
    "    model.compile(loss={'output_similarity': 'binary_crossentropy'},\n",
    "                  optimizer=Adam(1e-5),\n",
    "                  metrics={'output_similarity': 'accuracy'})\n",
    "\n",
    "    return model, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model, tokenizer = construct_model(paths)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build DataSequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QuestionCondPairsDataseq(Sequence):\n",
    "    def __init__(self, dataset, tokenizer, is_train=True, max_len=120, \n",
    "                 sampler=None, shuffle=False, batch_size=32):\n",
    "        self.dataset = dataset\n",
    "        self.tokenizer = tokenizer\n",
    "        self.is_train = is_train\n",
    "        self.max_len = max_len\n",
    "        self.sampler = sampler\n",
    "        self.shuffle = shuffle\n",
    "        self.batch_size = batch_size\n",
    "        self.on_epoch_end()       \n",
    "    \n",
    "    def _pad_sequences(self, seqs, max_len=None):\n",
    "        return pad_sequences(seqs, maxlen=max_len, padding='post', truncating='post')\n",
    "    \n",
    "    def __getitem__(self, batch_id):\n",
    "        batch_data_indices = \\\n",
    "            self.global_indices[batch_id * self.batch_size: (batch_id + 1) * self.batch_size]\n",
    "        batch_data = [self.data[i] for i in batch_data_indices]\n",
    "\n",
    "        X1, X2 = [], []\n",
    "        Y = []\n",
    "        \n",
    "        for data in batch_data:\n",
    "            x1, x2 = self.tokenizer.encode(first=data.question.lower(), \n",
    "                                           second=data.cond_text.lower())\n",
    "            X1.append(x1)\n",
    "            X2.append(x2)\n",
    "            if self.is_train:\n",
    "                Y.append([data.label])\n",
    "    \n",
    "        X1 = self._pad_sequences(X1, max_len=self.max_len)\n",
    "        X2 = self._pad_sequences(X2, max_len=self.max_len)\n",
    "        inputs = {'input_x1': X1, 'input_x2': X2}\n",
    "        if self.is_train:\n",
    "            Y = self._pad_sequences(Y, max_len=1)\n",
    "            outputs = {'output_similarity': Y}\n",
    "            return inputs, outputs\n",
    "        else:\n",
    "            return inputs\n",
    "                    \n",
    "    def on_epoch_end(self):\n",
    "        self.data = self.sampler.sample(self.dataset)\n",
    "        self.global_indices = np.arange(len(self.data))\n",
    "        if self.shuffle:\n",
    "            np.random.shuffle(self.global_indices)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return math.ceil(len(self.data) / self.batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_qc_pairs_seq = QuestionCondPairsDataseq(tr_qc_pairs, tokenizer, \n",
    "                                           sampler=NegativeSampler(), shuffle=True)\n",
    "\n",
    "te_qc_pairs_seq = QuestionCondPairsDataseq(te_qc_pairs, tokenizer, \n",
    "                                           sampler=FullSampler(), shuffle=False, batch_size=128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit_generator(tr_qc_pairs_seq, epochs=5, workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "te_result = model.predict_generator(te_qc_pairs_seq, verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make prediction for task2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_result(qc_pairs, result, threshold):\n",
    "    select_result = defaultdict(set)\n",
    "    for pair, score in zip(qc_pairs, result):\n",
    "        if score > threshold:\n",
    "            select_result[pair.query_id].update([pair.cond_sql])\n",
    "    return dict(select_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task2_result = merge_result(te_qc_pairs, te_result, threshold=0.995)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "final_output_file = 'final_output.json'\n",
    "with open(final_output_file, 'w') as f:\n",
    "    for query_id, pred_sql in enumerate(task1_result):\n",
    "        cond = list(task2_result.get(query_id, []))\n",
    "        pred_sql['conds'] = cond\n",
    "        json_str = json.dumps(pred_sql, ensure_ascii=False)\n",
    "        f.write(json_str + '\\n')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
